import tensorflow as tf
import os

tf.compat.v1.app.flags.DEFINE_string("ps_hosts", "192.168.2.43:2223,192.168.2.53:2223", "ps hosts")
tf.compat.v1.app.flags.DEFINE_string("worker_hosts", "192.168.2.43:2222,192.168.2.53:2222", "worker hosts")
tf.compat.v1.app.flags.DEFINE_string("job_name", "worker", "'ps' or'worker'")
tf.compat.v1.app.flags.DEFINE_integer("task_index", 0, "Index of task within the job")

FLAGS = tf.compat.v1.app.flags.FLAGS

def main(_):
    if FLAGS.job_name=='ps':
        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = '1'
        
    ps_hosts = FLAGS.ps_hosts.split(',')
    print(f"ps_hosts={ps_hosts}")
    worker_hosts = FLAGS.worker_hosts.split(',')
    cluster = tf.compat.v1.train.ClusterSpec(
        {
            "ps":ps_hosts,
            "worker":worker_hosts
        }
    )
    server = tf.compat.v1.train.Server(
        cluster,
        job_name=FLAGS.job_name,
        task_index = FLAGS.task_index
    )
    server.join()
    
if __name__ == '__main__':
    tf.compat.v1.app.run()